#include "vio/front_end/camera.h"

namespace vio {
Camera::Camera() {}

cv::Point3f Camera::KB_unprojection(const cv::Point2f& p2) {
  Eigen::Vector3d p3_cam = K_in_inv_ * Eigen::Vector3d(p2.x, p2.y, 1.0);
  return cv::Point3f(p3_cam[0], p3_cam[1], p3_cam[2]);
}
cv::Point2f Camera::KB_projection(const cv::Point3f& p3_cam) {
  Eigen::Vector3d p3_cam_e(p3_cam.x, p3_cam.y, p3_cam.z);
  Eigen::Vector2d p2_e = KB_projection(p3_cam_e);
  return cv::Point2f(p2_e[0], p2_e[1]);
}

Eigen::Vector3d Camera::KB_unprojection(const Eigen::Vector3d& p2) {
  Eigen::Vector3d p3_cam = K_in_inv_ * p2;
  return p3_cam;
}

Eigen::Vector2d Camera::KB_projection(const Eigen::Vector3d& p3_cam) {
  double theta = std::acos(p3_cam[2] / p3_cam.norm());
  double phi = std::atan2(p3_cam[1], p3_cam[0]);
  // D： distortion params
  double k2 = D_[0], k3 = D_[1], k4 = D_[2], k5 = D_[3];
  // K: converts back to pixel coordinates in the output image
  double fx = K_out_(0, 0), cx = K_out_(0, 2), fy = K_out_(1, 1),
         cy = K_out_(1, 2);
  double theta2 = theta * theta;
  double theta3 = theta2 * theta;
  double theta5 = theta2 * theta3;
  double theta7 = theta5 * theta2;
  double theta9 = theta7 * theta2;
  double d_theta =
      theta + k2 * theta3 + k3 * theta5 + k4 * theta7 + k5 * theta9;
  Eigen::Vector2d p_u(std::cos(phi) * d_theta, std::sin(phi) * d_theta);
  Eigen::Vector2d pp;
  pp[0] = fx * p_u[0] + cx;
  pp[1] = fy * p_u[1] + cy;
  return pp;
}

void Camera::undistortImage(const cv::Mat& img, cv::Mat& out_img) {
  int rows = img.rows;
  int cols = img.cols;

  cv::Mat mapX = cv::Mat::zeros(rows, cols, CV_32F);
  cv::Mat mapY = cv::Mat::zeros(rows, cols, CV_32F);
  for (int v = 0; v < rows; v++) {
    for (int u = 0; u < cols; u++) {
      Eigen::Vector3d p3_cam = KB_unprojection(Eigen::Vector3d(u, v, 1));
      Eigen::Vector2d p2_pixel = KB_projection(p3_cam);
      mapX.at<float>(v, u) = p2_pixel(0);
      mapY.at<float>(v, u) = p2_pixel(1);
    }
  }
  cv::Mat map1, map2;
  cv::convertMaps(mapX, mapY, map1, map2, CV_32FC1, false);
  cv::remap(img, out_img, map1, map2, cv::INTER_LINEAR);
}

Eigen::Vector3d Camera::world2camera(const Eigen::Vector3d& p_w,
                                     const Sophus::SE3d& T_c_w) {
  return T_c_w * p_w;
}

Eigen::Vector3d Camera::camera2world(const Eigen::Vector3d& p_c,
                                     const Sophus::SE3d& T_c_w) {
  return T_c_w.inverse() * p_c;
}

Eigen::Vector2d Camera::camera2pixel(const Eigen::Vector3d& p_c) {
  return Eigen::Vector2d(fx_ * p_c(0, 0) / p_c(2, 0) + cx_,
                         fy_ * p_c(1, 0) / p_c(2, 0) + cy_);
}

cv::Point2f Camera::camera2pixel(const cv::Point3f& p_c) {
  return cv::Point2f(fx_ * p_c.x / p_c.z + cx_, fy_ * p_c.y / p_c.z + cy_);
}

Eigen::Vector3d Camera::pixel2camera(const Eigen::Vector2d& p_p, double depth) {
  return Eigen::Vector3d(
      (p_p(0, 0) - cx_) * depth / fx_, (p_p(1, 0) - cy_) * depth / fy_, depth);
}

cv::Point3f Camera::pixel2camera(const cv::Point2f& p_p, double depth) {
  return cv::Point3f(
      (p_p.x - cx_) * depth / fx_, (p_p.y - cy_) * depth / fy_, depth);
}

Eigen::Vector2d Camera::world2pixel(const Eigen::Vector3d& p_w,
                                    const Sophus::SE3d& T_c_w) {
  return camera2pixel(world2camera(p_w, T_c_w));
}

Eigen::Vector3d Camera::pixel2world(const Eigen::Vector2d& p_p,
                                    const Sophus::SE3d& T_c_w,
                                    double depth) {
  return camera2world(pixel2camera(p_p, depth), T_c_w);
}
}  // namespace vio